import argparse

from load_data import dataset
from GMM import GMM
from kmeans import Kmeans

"""
    Main Function for GMM (Gaussian Mixture Mode)
    ---------------------------------------------

    default variable:
    - n: Number of Clusters = 4
    - i: Number of Iteration = 50
    - inter: Interval between Every Frame = 0.1
    - model: Modle you want to use = 'GMM' choices = ['GMM', 'kmeans']
    
    Note: init only for GMM 
    - init: Initialization Mode = 'random' choices = ['random', 'kmeans']
    ---------------------------------------------
    
    Note: 
    If you want to change the default variable,
    please check following dict (_default_variable), 
    and change the variable you want to modify.
    ---------------------------------------------
"""

_default_variable = {
    'n': 3,  # choices = [3, 4, 6, 8]
    'i': 50,
    'inter': 0.1,
    'init': 'random',  # random or kmeans
    'model': 'GMM'  # GMM or kmeans
}


def _run(num_class, max_iter, init, interval=0.1, model='GMM'):
    """
    Call GMM Model
    :param num_class:  number of class/cluster
    :param max_iter: number of iteration
    :param init: initialization mode
    :param interval: refreshing interval
    :return:
    """
    _, X = dataset.load_data(num_class)
    if model == 'GMM':
        model = GMM(X, num_class, max_iter=max_iter, interval=interval)
        model.fit(init=init, visual=True)
    elif model == 'kmeans':
        model = Kmeans(X, num_class)
        model.fit()
    else:
        raise ValueError("I'm sorry. I only implement two algorithms, GMM and kmeans."
                         "So, you have to choose one from them")


def _argument_parser():
    """
    Parse argument form CLI
    -------
    default values:
        num_cluster = 4
        max_iter = 50
        init = 'random'
        interval = 0.1

    :return: args
    """
    parser = argparse.ArgumentParser(description='Main Function for GMM Algorithm')
    parser.add_argument('-n', dest='num_cluster', default=_default_variable['n'],
                        metavar='Number of Clusters: 4', type=int,
                        choices=[3, 4, 6, 8])

    parser.add_argument('-i', metavar='Number of Iteration: 50', type=int,
                        dest='max_iter', default=_default_variable['i'])

    parser.add_argument('-inter', metavar='Interval between Every Frame: 0.1', type=float,
                        dest='interval', default=_default_variable['inter'])

    parser.add_argument('-init', metavar='Initialization Mode : random', type=str,
                        dest='init', default=_default_variable['init'], choices=['random', 'kmeans'])
    parser.add_argument('-model', metavar='Model you want to use', type=str,
                        dest='model', default='GMM', choices=['GMM', 'kmeans'])
    args = parser.parse_args()

    return args


def main():
    """
    main function for GMM algorithm
    :return: None

    -------
    default values:
        num_cluster = 4
        max_iter = 50
        init = 'random'
        interval = 0.1
    """
    args = _argument_parser()
    _run(args.num_cluster, args.max_iter, args.init, interval=args.interval, model=args.model)


if __name__ == '__main__':
    main()
